
# coding: utf-8

# In[ ]:

"""Part 1:Read and parse the initial dataset"""

"""(1a) Load and check the data"""

#load testing library
from test_helper import Test
import os.path
baseDir = os.path.join('data')
inputPath = os.path.join('#', '#')
fileName = os.path.join(baseDir, inputPath)

numPartitions = 2
rawData = sc.textFile(fileName, numPartitions)


numPoints = rawData.count()
print numPoints
samplePoints = rawData.take(5)
print samplePoints

numPoints = rawData.count()
print numPoints
samplePoints = rawData.take(5)
print samplePoints


# In[ ]:

"""(1b) Using LabeledPoint"""

# Here is a sample raw data point:
# '2001.0,0.884,0.610,0.600,0.474,0.247,0.357,0.344,0.33,0.600,0.425,0.60,0.419'
# In this raw data point, 2001.0 is the label, and the remaining values are features

from pyspark.mllib.regression import LabeledPoint
import numpy as up

def parsePoint(line):
    """Converts a comma separated unicode string into a `LabeledPoint`.

    Args:
        line (unicode): Comma separated unicode string where the first element is the label and the
            remaining elements are features.

    Returns:
        LabeledPoint: The line is converted into a `LabeledPoint`, which consists of a label and
            features.
    """
    lines = line.split(',')
    
    return LabeledPoint(lines[0], lines[1:])

parsedSamplePoints = map(parsePoint, samplePoints)
firstPointFeatures = parsedSamplePoints[0].features
firstPointLabel = parsedSamplePoints[0].label
print firstPointFeatures, firstPointLabel

d = len(firstPointFeatures)
print d


# In[ ]:

"""Visualization 1: Features"""

"""
   First we will load and setup the visualization library. Then we will look at the raw features for 50 data points by generating a heatmap 
   that visualizes each feature on a grey-scale and shows the variation of each feature across the 50 sample data points. The features are 
   all between 0 and 1, with values closer to 1 represented via darker shades of grey.
"""
import matplotlib.pyplot as plt
import matplotlib.cm as cm

sampleMorePoints = rawData.take(50)
# You can uncomment the line below to see randomly selected features.  These will be randomly
# selected each time you run the cell.  Note that you should run this cell with the line commented
# out when answering the lab quiz questions.
# sampleMorePoints = rawData.takeSample(False, 50)

parsedSampleMorePoints = map(parsePoint, sampleMorePoints)
dataValues = map(lambda lp: lp.features.toArray(), parsedSampleMorePoints)

def preparePlot(xticks, yticks, figsize=(10.5, 6), hideLabels=False, gridColor='#999999',
                gridWidth=1.0):
    """Template for generating the plot layout."""
    plt.close()
    fig, ax = plt.subplots(figsize=figsize, facecolor='white', edgecolor='white')
    ax.axes.tick_params(labelcolor='#999999', labelsize='10')
    for axis, ticks in [(ax.get_xaxis(), xticks), (ax.get_yaxis(), yticks)]:
        axis.set_ticks_position('none')
        axis.set_ticks(ticks)
        axis.label.set_color('#999999')
        if hideLabels: axis.set_ticklabels([])
    plt.grid(color=gridColor, linewidth=gridWidth, linestyle='-')
    map(lambda position: ax.spines[position].set_visible(False), ['bottom', 'top', 'left', 'right'])
    return fig, ax

# generate layout and plot
fig, ax = preparePlot(np.arange(.5, 11, 1), np.arange(.5, 49, 1), figsize=(8,7), hideLabels=True,
                      gridColor='#eeeeee', gridWidth=1.1)
image = plt.imshow(dataValues,interpolation='nearest', aspect='auto', cmap=cm.Greys)
for x, y, s in zip(np.arange(-.125, 12, 1), np.repeat(-.75, 12), [str(x) for x in range(12)]):
    plt.text(x, y, s, color='#999999', size='10')
plt.text(4.7, -3, 'Feature', color='#999999', size='11'), ax.set_ylabel('Observation')
pass


# In[ ]:

"""(1c) Find the range"""

parsedDataInit = rawData.map(parsePoint)
onlyLabels = parsedDataInit.map(lambda x:x.label)
minYear = onlyLabels.min()
maxYear = onlyLabels.max()
print maxYear, minYear


# In[ ]:

"""(1d) Shift labels"""
"""
   shifting labels such that they start from zero
"""
parsedData = parsedDataInit.map(lambda x:LabeledPoint(x.label-minYear, x.features))

#Should be a LabeledPoint
print type(parsedData.take(1)[0])


# In[ ]:

"""Visualization 2: Shifting labels"""

"""
   We will look at the labels before and after shifting them. Both scatter plots below visualize tuples storing i) a label value and ii) the 
   number of training points with this label. The first scatter plot uses the initial labels, while the second one uses the shifted labels.
"""
#get data for plot
oldData = (parsedDataInit
          .map(lambda lp:(lp.label, 1))
          .reduceByKey(lambda x, y:x+y)
          .collect())
x, y = zip(*oldData)

#generate layout and plot data
fig, ax = preparePlot(np.arange(1920, 2050, 20), np.arange(0, 150, 20))
plt.scatter(x, y, s=14**2, c='#d6ebf2', edgecolors='#8cbfd0', alpha=0.75)
ax.set_xlabel('Year'), ax.set_ylabel('Count')
pass


# In[ ]:

"""(1e) Training validation and test sets"""
weights = [.8, .1, .1]
seed = 42
parsedTrainData, parsedValData, parsedTestData = parsedData.randomSplit(weights, seed)
parsedTraindata.persist
parsedValData.persist
parsedTestData.persist
nTrain = parsedTrainData.count()
nVal = parsedValData.count()
nTest = parsedTestData.count()

print nTrain, nVal, nTest, nTrain + nVal + nTest
print parsedData.count()


# In[ ]:

"""Part 2: Create and evaluate a baseline model"""

"""(2a) Average label"""
"""
   A very simple yet natural baseline model is one where we always make the same prediction independent of the given data point, using the average
   label in the training set as the constant perdiction value. Compute this value, which is the vaerage (shifted) song year for the training set.
"""
averageTrainYear = (parsedTrainData
                   .map(lambda x:x.label).reduce(lambda x, y:x+y))/nTrain

print averageTrainYear


# In[ ]:

"""(2b) Root mean squared error"""
"""
   We naturally would like to see how well this naive baseline performs. We will use root mean squared error(RMSE) for evaluation purposes. Implement
   a function to compute RMSE given an RDD of (label, prediction) tuples, and test out this function on an example.
"""
import math
def squaredError(label, prediction):
    """Calculates the squared error for a single prediction.
    
    Args:
        label (float): The correct value for this observation.
        prediction(float): The predicted value for this observation.
        
    Returns:
        float: The difference between the label and prediction squared.
    """
    return math.pow((label-prediction), 2)

def calcRMSE(labelsAndPreds):
    """Calculates the root mean squared error for an RDD of(label, predition) tuples.
    
    Args:
        labelsAndPred (RDD of(float, float)): An 'RDD' of (label, prediction) tuples.
        
    Returns:
        float: The square root of the mean of the squared errors.
    """
    meanSquaredError = (labelsAndPreds.map(lambda (x,y):squaredError(x,y)).reduce(lambda x, y:x+y))/labelsAndPreds.count()
    return math.sqrt(meanSquaredError)

#Test Root mean squared error
Test.assertTrue(np.allclose(squaredError(3, 1), 4.)'incorrect definition of squaredError')
Test.assertTrue(np.allclose(exampleRMSE, 1.29099444874), 'incorrect value for exampleRMSE')


# In[ ]:

"""(2c) Training, validation and test RMSE"""
"""
   Now let's calcuate the training, validation and test RMSE of our baseline model.
"""

labelsAndPredsTrain = parsedTrainData.parsedTrainData.map(lambda x:(x.label, averageTrainYear))
rmseTrainBase = calcRMSE(labelsAndPredsTrain)

labelsAndPredsVal = parsedValData.map(lambda x:(x.label, averageTrainYear))
rmseValBase = calcRMSE(labelsAndPredsVal)

labelsAndPredsTest = parsedTestData.map(lambda x:(x.label, averageTrainYear))
rmseTestBase = calcRMSE(labelsAndPredsTest)

print 'Baseline Train RMSE = {0:.3f}'.format(rmseTrainBase)
print 'Baseline validation RMSE = {0:.3f}'.format(rmseValBase)
print 'Baseline Test RMSE = {0:.3f}'.format(rmseTestBase)


# In[ ]:

"""Visualization 3: Predicted vs. actual"""
"""
   We will visualize predictions on the validation dataset. The scatter plots below visualize tuples storing i) the predicted value and ii) true
   label. The first scatter plot represents the ideal situation where the predicted value exactly equals the true label, while the second plot 
   uses the baseline predictor (i.e., averageTrainYear) for all predicted values. Further note that the points in the scatter plots are color-coded,
   ranging from light yellow when the true and predicted values are equal to bright red when they drastically differ.
"""
from matplotlib.colors import ListedColormap, Normalize
from matplotlib.cm import get_cmap
cmap = get_cmap('YlOrRd')
norm = Normalize()

actual = np.asarray(parsedValData
                    .map(lambda lp: lp.label)
                    .collect())
error = np.asarray(parsedValData
                   .map(lambda lp: (lp.label, lp.label))
                   .map(lambda (l, p): squaredError(l, p))
                   .collect())
clrs = cmap(np.asarray(norm(error)))[:,0:3]
fig, ax = preparePlot(np.arange(0, 100, 20), np.arange(0, 100, 20))
plt.scatter(actual, actual, s=14**2, c=clrs, edgecolors='#888888', alpha=0.75, linewidths=0.5)
ax.set_xlabel('Predicted'), ax.set_ylabel('Actual')
pass


# In[ ]:

predictions = np.asarray(parsedValData
                         .map(lambda lp: averageTrainYear)
                         .collect())
error = np.asarray(parsedValData
                   .map(lambda lp: (lp.label, averageTrainYear))
                   .map(lambda (l, p): squaredError(l, p))
                   .collect())
norm = Normalize()
clrs = cmap(np.asarray(norm(error)))[:,0:3]

fig, ax = preparePlot(np.arange(53.0, 55.0, 0.5), np.arange(0, 100, 20))
ax.set_xlim(53, 55)
plt.scatter(predictions, actual, s=14**2, c=clrs, edgecolors='#888888', alpha=0.75, linewidths=0.3)
ax.set_xlabel('Predicted'), ax.set_ylabel('Actual')


# In[ ]:

"""Part 3: Train(via grandient descent) and evaluate a linear regression model"""

"""(3a) Gradient summand"""

"""
   Now let's see if we can do better via linear regression, training a model via gradient descent (we'll omit the intercept for now). Recall 
   that the gradient descent update for linear regression is:
                                               wi+1=wi−αi∑j(w⊤ixj−yj)xj.
   where i is the iteration number of the gradient descent algorithm, and j identifies the observation.  
   First, implement a function that computes the summand for this update, i.e., the summand equals (w⊤x−y)x, and test out this function on 
   two examples.
"""
from pyspark.mllib.linalg import DenseVector

def gradientSummand(weights, lp):
    """Calculates the gradient summand for a given weight and LabeledPoint.
    
    Args:
        weights(DenseVector): An array of model weights(beats).
        lp(LabeledPoint): The LabeledPoint for a single observation.
    
    Returns:
        DensenVector: An array of values the same length as weights. The gradient summand.
    """
    return (DensenVector.dot(weights, lp.features)-lp.label)*lp.features

#Test Gradient summand
exampleW = DenseVector([1, 1, 1])
exampleLP = LabeledPoint(2.0, [3, 1, 4])

summandOne = gradientSummand(exampleW, exampleLP)
print summandOne

exampleW = DenseVector([.24, 1.2, -1.4])
exampleLP = LabeledPoint(3.0, [-1.4, 4.2, 2.1])
summandTwo = gradientSummand(exampleW, exampleLP)
print summandTwo

Test.assertTrue(np.allclose(summandOne, [18., 6., 24.]), 'incorrect value for summandOne')
Test.assertTrue(np.allclose(summandTwo, [1.7304,-5.1912,-2.5956]), 'incorrect value for summandTwo')

    
    
    
    


# In[ ]:

"""(3b) Use weights to make predictions"""

def getLabeledPrediction(weights, observation):
    """Calculate predictions and returns a (label, prediction) tuple.
    
    Args:
        weights(np.ndarray): An array with one weight for each features in trainData.
        observation (LabeledPoint): A LabeledPoint that contain the correct label and the features for the data point.
        
    Returns:
        tuple: A(label, prediction) tuple.
    """
    return (observation.label, np.dot(weights, observation.features))

#Test Use weights to make predictions
weights = np.array([1.0, 1.5])
predictionExample = sc.parallelize([LabeledPoint(2, np.array([1.0, .5])),
                                   LabeledPoint(1.5, np.array([.5, .5]))])
labelsAndPredsExample = predictionExample.map(lambda lp: getLabeledPrediction(weights, lp))
print labelsAndPredsExample.collect()

Test.assertEquals(labelsAndPredsExample.collect(), [(2.0, 1.75), (1.5, 1.25)],
                  'incorrect definition for getLabeledPredictions')


# In[ ]:

"""(3c) Gradient descent"""

def linregGradientDescent(trainData, numIters):
    """Calculates the weights and error for a linear regression model trained with gradient descent.
    
    Args:
        trainData (RDD of LabeledPoint): The labeled data for use in training the model.
        numIters (int): The number of iterations of gradient descent to perform.
        
    Returns:
        (np.ndarray, np.ndarray): A tuple of (weights, training errors).  Weights will be the
            final weights (one weight per feature) for the model, and training errors will contain
            an error (RMSE) for each iteration of the algorithm.
    """
    #The length of the training data
    n = trainData.count()
    #The number of features in the training data
    d = len(trainData.take(1)[0].features)
    w = np.zeros(d)
    alpha = 1.0
    #We will compute and store the training error after each iteration
    errorTrain = np.zeros(numIters)
    for i in range(numIters):
        #Use getLabeledPrediction with trainData to obtain an RDD of (label, prediction) tuples.Note that the weights all equal 0 for the first
        #iteration, so the predictions will have large errors to start.
        labelsAndPredsTrain = trainData.map(lambda x:getLabeledPrediction(w, x))
        errorTrain[i] = calcRMSE(labelsAndPredsTrain)
        
        #Calculate the gradient. Make use of the gradientSummand function 
        #Note that gradient should be a DenseVector of length d
        w = DenseVector(w)
        gradient = trainData.map(lambda x:gradientSummand(w, x)).sum()
        
        #Update the weights
        alpha_i = alpha/(n*np.sqrt(i+1))
        w -= alpha_i * gradient
    return w, errorTrain

#create a toy dataset with n=10, d=3, and then run 5 iterators of gradient descent to verify that linregGradientDescent is working properly
exampleN = 10
exampleD = 3
exampleData = (sc
               .parallelize(parsedTrainData.take(exampleN))
               .map(lambda lp: LabeledPoint(lp.label, lp.features[0:exampleD])))
print exampleData.take(2)
exampleNumIters = 5
exampleWeights, exampleErrorTrain = linregGradientDescent(exampleData, exampleNumIters)
print exampleWeights


# In[ ]:

"""(3d) Train the model"""
"""
   Now let's train a linear regression model on all of our training data and evaluate its accuracy on the validation set.
"""
numIters = 50
weightsLRO, errorTrainLRO = linregGradientDescent(parsedTrainData, numIters)

labelsAndPreds = parsedValData.map(lambda x:(x.label, averageTrainYear))
rmseValLRO = calcRMSE(labelsAndPreds)

print 'Validation RMSE:\n\tBaseline = {0:.3f}\n\tLR0 = {1:.3f}'.format(rmseValBase,
                                                                       rmseValLR0) 


# In[ ]:

"""Visualization 4: Training error"""
"""
   We will look at the log of the training error as a function of iteration. The first scatter plot visualizes the logarithm of the training error
   for all 50 iterations. The second plot shows the training error itself, focusing on the final 44 iterations.
"""
norm = Normalize()
clrs = cmap(np.asarray(norm(np.log(errorTrainLR0))))[:,0:3]

fig, ax = preparePlot(np.arange(0, 60, 10), np.arange(2, 6, 1))
ax.set_ylim(2, 6)
plt.scatter(range(0, numIters), np.log(errorTrainLR0), s=14**2, c=clrs, edgecolors='#888888', alpha=0.75)
ax.set_xlabel('Iteration'), ax.set_ylabel(r'$\log_e(errorTrainLR0)$')
pass


# In[ ]:

norm = Normalize()
clrs = cmap(np.asarray(norm(errorTrainLR0[6:])))[:,0:3]

fig, ax = preparePlot(np.arange(0, 60, 10), np.arange(17, 22, 1))
ax.set_ylim(17.8, 21.2)
plt.scatter(range(0, numIters-6), errorTrainLR0[6:], s=14**2, c=clrs, edgecolors='#888888', alpha=0.75)
ax.set_xticklabels(map(str, range(6, 66, 10)))
ax.set_xlabel('Iteration'), ax.set_ylabel(r'Training Error')
pass


# In[ ]:

"""Part 4: Train using MLlib and perform grid search"""

"""(4a) LinearRegressionWithSGD"""

"""
   We are already doing better than the baseline model, but let's see if we can do better by adding an intercept, using regularization, and training 
   for more iterations. MLlib's LinearRegressionWithSGD essentially implements the same algorithm that we implemented in Part(3b)
"""
from pyspark.mllib.regression import LinearRegressionWithSGD
#Values to use when training the linear regression model
numIters = 500 #iterations
alpha = 1.0 #step
miniBatchFrac = 1.0 #miniBatchFraction
reg = 1e-1 #regParam
regType = '12' #regType
useIntercept = True #intercept

firstModel = LinearRegressionWithSGD.train(parsedTrainData, iterations = 500, step = 1.0, miniBatchFraction = 1.0, regParam = 1e-1, regType = '12',
                                          intercept = True)

#weightsLR1 stores the model weights, interceptLR1 stores the model intercept
weightsLR1 = firstModel.weights
interceptLR1 = firstModel.intercept
print weightsLR1, interceptLR1


# In[ ]:

"""(4b) Predict"""
"""
   Using the LinearRegressionModel.predict() method to make a prediction on a sample point.
"""
samplePoint = parsedTrainData.take(1)[0]
samplePrediction = firstModel.predict(samplePoint.features)
print samplePrediction


# In[ ]:

"""(4c) Evaluate RMSE"""
"""
   Evaluating the accuracy of this model on the validation set.
"""
labelsAndPreds = parsedValData.map(lambda lp: (lp.label, firstModel.predict(lp.features)))

rmseValLR1 = calcRMSE(labelsAndPreds)

print ('Validation RMSE:\n\tBaseline = {0:.3f}\n\tLRO = {1:.3f}' +
      '\n\tLR1 = {2:.3f}').format(rmseValBase, rmseValLRO, rmseValLR1)


# In[ ]:

"""(4d) Grid search"""
"""
   Performing grid search to find a good regularization parameter. Try regParam values 1e-10, 1e-5, and 1
"""
bestRMSE = rmseValLR1
bestRegParam = reg
bestModel = firstModel

numIters = 500
alpha = 1.0
miniBatchFrac = 1.0
for reg in [1e-10, 1e-5, 1]:
    model = LinearRegressionWithSGD.train(parsedTrainData, numIters, alpha, miniBatchFrac, regParam = reg,
                                         regType = '12', intercept = True)
    labelsAndPreds = parsedValData.map(lambda lp:(lp.label, model.predict(lp.features)))
    rmseValGrid = calcRMSE(labelsAndPreds)
    print rmseValGrid
    
    if rmseValGrid < bestRMSE
        bestRMSE = rmseValGrid
        bestRegParam = reg
        bestModel = model

rmseValLRGrid = bestRMSE

print ('Validation RMSE:\n\tBaseline = {0:.3f}\n\tLR0 = {1:.3f}\n\tLR1 = {2:.3f}\n' +
       '\tLRGrid = {3:.3f}').format(rmseValBase, rmseValLR0, rmseValLR1, rmseValLRGrid)


# In[ ]:

"""Visualization 5: Best model's predictions"""
"""
   We create a color-coded scatter plot visualizing tuples storing i) the predicted value from this model and ii) true label.
"""
predictions = np.asarray(parsedValData
                         .map(lambda lp: bestModel.predict(lp.features))
                         .collect())
actual = np.asarray(parsedValData
                    .map(lambda lp: lp.label)
                    .collect())
error = np.asarray(parsedValData
                   .map(lambda lp: (lp.label, bestModel.predict(lp.features)))
                   .map(lambda (l, p): squaredError(l, p))
                   .collect())

norm = Normalize()
clrs = cmap(np.asarray(norm(error)))[:,0:3]

fig, ax = preparePlot(np.arange(0, 120, 20), np.arange(0, 120, 20))
ax.set_xlim(15, 82), ax.set_ylim(-5, 105)
plt.scatter(predictions, actual, s=14**2, c=clrs, edgecolors='#888888', alpha=0.75, linewidths=.5)
ax.set_xlabel('Predicted'), ax.set_ylabel(r'Actual')
pass


# In[ ]:

"""(4e) Vary alpha and the number of iterations"""
"""
   try 1e-5 and 10 as values for alpha and also try training models for 500 iterations (as before) but also for 5 iterations. Evaluate all 
   models on the validation set.
"""
reg = bestRegParam
modelRMSEs = []

for alpha in [1e-5, 10]:
    for numIters in [500, 5]:
        model = LinearRegressionWithSGD.train(parsedTrainData, numIters, alpha, miniBatchFrac, regParam = reg,
                                             regType = '12', intercept = True)
        labelsAndPreds = parsedValData.map(lambda lp: (lp.label, model.predict(lp.features)))
        rmseVal = calcRMSE(labelsAndPreds)
        print 'alpha = {0:.0e}, numIters = {1}, RMSE = {2:.3f}'.format(alpha, numIters, rmseVal)
        modelRMSEs.append(rmseVal)


# In[ ]:

"""Visualization 6: Hyperparameter heat map"""
"""
   We create a heat map where the brighter colors correspond to lower RMSE values. The first plot has a large area with brighter colors. In 
   order to differentiate within the bright region, we generate a second plot corresponding to the hyperparameters found within that region.
"""
from matplotlib.colors import LinearSegmentedColormap

# Saved parameters and results, to save the time required to run 36 models
numItersParams = [10, 50, 100, 250, 500, 1000]
regParams = [1e-8, 1e-6, 1e-4, 1e-2, 1e-1, 1]
rmseVal = np.array([[  20.36769649,   20.36770128,   20.36818057,   20.41795354,  21.09778437,  301.54258421],
                    [  19.04948826,   19.0495    ,   19.05067418,   19.16517726,  19.97967727,   23.80077467],
                    [  18.40149024,   18.40150998,   18.40348326,   18.59457491,  19.82155716,   23.80077467],
                    [  17.5609346 ,   17.56096749,   17.56425511,   17.88442127,  19.71577117,   23.80077467],
                    [  17.0171705 ,   17.01721288,   17.02145207,   17.44510574,  19.69124734,   23.80077467],
                    [  16.58074813,   16.58079874,   16.58586512,   17.11466904,  19.6860931 ,   23.80077467]])

numRows, numCols = len(numItersParams), len(regParams)
rmseVal = np.array(rmseVal)
rmseVal.shape = (numRows, numCols)

fig, ax = preparePlot(np.arange(0, numCols, 1), np.arange(0, numRows, 1), figsize=(8, 7), hideLabels=True,
                      gridWidth=0.)
ax.set_xticklabels(regParams), ax.set_yticklabels(numItersParams)
ax.set_xlabel('Regularization Parameter'), ax.set_ylabel('Number of Iterations')

colors = LinearSegmentedColormap.from_list('blue', ['#0022ff', '#000055'], gamma=.2)
image = plt.imshow(rmseVal,interpolation='nearest', aspect='auto',
                    cmap = colors)


# In[ ]:

# Zoom into the bottom left
numItersParamsZoom, regParamsZoom = numItersParams[-3:], regParams[:4]
rmseValZoom = rmseVal[-3:, :4]

numRows, numCols = len(numItersParamsZoom), len(regParamsZoom)

fig, ax = preparePlot(np.arange(0, numCols, 1), np.arange(0, numRows, 1), figsize=(8, 7), hideLabels=True,
                      gridWidth=0.)
ax.set_xticklabels(regParamsZoom), ax.set_yticklabels(numItersParamsZoom)
ax.set_xlabel('Regularization Parameter'), ax.set_ylabel('Number of Iterations')

colors = LinearSegmentedColormap.from_list('blue', ['#0022ff', '#000055'], gamma=.2)
image = plt.imshow(rmseValZoom,interpolation='nearest', aspect='auto',
                    cmap = colors)
pass


# In[ ]:

"""Part 5: Add interactions between features"""

"""(5a) Add 2-way interactions"""
"""
   We're used the features as they were provided. Now, we will add features that capture the two-way interactions between our existion features.
"""
import itertools

def twoWayInteractons(lp):
    """Creates a new `LabeledPoint` that includes two-way interactions.

    Note:
        For features [x, y] the two-way interactions would be [x^2, x*y, y*x, y^2] and these
        would be appended to the original [x, y] feature list.

    Args:
        lp (LabeledPoint): The label and features for this observation.

    Returns:
        LabeledPoint: The new `LabeledPoint` should have the same label as `lp`.  Its features
            should include the features from `lp` followed by the two-way interaction features.
    """
    interactions = map(lambda (x, y):x*y, itertools.product(lp.features, repeat = 2))
    return LabeledPoint(lp.label, np.hstack((lp.features, interactions)))

print twoWayInteractions(LabeledPoint(0.0, [2, 3]))

#Transform the existing train, validation, and test sets to include two-way interactions.
trainDataInteract = parsedTrainData.map(twoWayInteractions)
valDataInteract = parsedValData.map(twoWayInteractions)
testDataInteract = parseTestData.map(twoWayInteractions)

#Test Add two-way interactions
twoWayExample = twoWayInteractions(LabeledPoint(0.0, [2, 3]))
Test.assertTrue(np.allclose(sorted(twoWayExample.features),
                            sorted([2.0, 3.0, 4.0, 6.0, 6.0, 9.0])),
                'incorrect features generatedBy twoWayInteractions')
twoWayPoint = twoWayInteractions(LabeledPoint(1.0, [1, 2, 3]))
Test.assertTrue(np.allclose(sorted(twoWayPoint.features),
                            sorted([1.0,2.0,3.0,1.0,2.0,3.0,2.0,4.0,6.0,3.0,6.0,9.0])),
                'incorrect features generated by twoWayInteractions')
Test.assertEquals(twoWayPoint.label, 1.0, 'incorrect label generated by twoWayInteractions')


# In[ ]:

"""(5b) Build interaction model"""

numIters = 500
alpha = 1.0
miniBatchFrac = 1.0
reg = 1e-10

modelInteract = LinearRegressionWithSGD.train(trainDataInteract, numIters, alpha, miniBatchFrac, regParam = reg,
                                             regType = '12', intercept = True)
labelsAndPredsInteract = valDataInteract.map(lambda lp: (lp.label, modelInteract.predict(lp.features)))

print ('Validation RMSE:\n\tBaseline = {0:.3f}\n\tLR0 = {1:.3f}\n\tLR1 = {2:.3f}\n\tLRGrid = ' +
       '{3:.3f}\n\tLRInteract = {4:.3f}').format(rmseValBase, rmseValLR0, rmseValLR1,
                                                 rmseValLRGrid, rmseValInteract)


# In[ ]:

"""(5c) Evaluate interaction model on test data"""

labelsAndPredsTest = testDataInteract.map(lambda lp: (lp.label, modelInteract.predict(lp.features)))

rmseTestInteract = calcRMSE(labelsAndPredsTest)

print ('Test RMSE:\n\tBaseline = {0:.3f}\n\tLRInteract = {1:.3f}'
       .format(rmseTestBase, rmseTestInteract))

